// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/exec/order_by.h>
#include <pollux/exec/operator_utils.h>
#include <pollux/exec/task.h>
#include <pollux/vector/flat_vector.h>

namespace kumo::pollux::exec {
    namespace {
        CompareFlags fromSortOrderToCompareFlags(const core::SortOrder &sortOrder) {
            return {
                sortOrder.isNullsFirst(),
                sortOrder.isAscending(),
                false,
                CompareFlags::NullHandlingMode::kNullAsValue
            };
        }
    } // namespace

    OrderBy::OrderBy(
        int32_t operatorId,
        DriverCtx *driverCtx,
        const std::shared_ptr<const core::OrderByNode> &orderByNode)
        : Operator(
            driverCtx,
            orderByNode->outputType(),
            operatorId,
            orderByNode->id(),
            "OrderBy",
            orderByNode->canSpill(driverCtx->queryConfig())
                ? driverCtx->makeSpillConfig(operatorId)
                : std::nullopt) {
        maxOutputRows_ = outputBatchRows(std::nullopt);
        POLLUX_CHECK(pool()->trackUsage());
        std::vector<column_index_t> sortColumnIndices;
        std::vector<CompareFlags> sortCompareFlags;
        sortColumnIndices.reserve(orderByNode->sortingKeys().size());
        sortCompareFlags.reserve(orderByNode->sortingKeys().size());
        for (int i = 0; i < orderByNode->sortingKeys().size(); ++i) {
            const auto channel =
                    exprToChannel(orderByNode->sortingKeys()[i].get(), outputType_);
            POLLUX_CHECK(
                channel != kConstantChannel,
                "OrderBy doesn't allow constant sorting keys");
            sortColumnIndices.push_back(channel);
            sortCompareFlags.push_back(
                fromSortOrderToCompareFlags(orderByNode->sortingOrders()[i]));
        }
        sortBuffer_ = std::make_unique<SortBuffer>(
            outputType_,
            sortColumnIndices,
            sortCompareFlags,
            pool(),
            &nonReclaimableSection_,
            driverCtx->prefixSortConfig(),
            spillConfig_.has_value() ? &(spillConfig_.value()) : nullptr,
            &spillStats_);
    }

    void OrderBy::addInput(RowVectorPtr input) {
        sortBuffer_->addInput(input);
    }

    void OrderBy::reclaim(
        uint64_t targetBytes,
        memory::MemoryReclaimer::Stats &stats) {
        POLLUX_CHECK(canReclaim());
        POLLUX_CHECK(!nonReclaimableSection_);

        // TODO: support fine-grain disk spilling based on 'targetBytes' after
        // having row container memory compaction support later.
        sortBuffer_->spill();

        // Release the minimum reserved memory.
        pool()->release();
    }

    void OrderBy::noMoreInput() {
        Operator::noMoreInput();
        sortBuffer_->noMoreInput();
        maxOutputRows_ = outputBatchRows(sortBuffer_->estimateOutputRowSize());
    }

    RowVectorPtr OrderBy::getOutput() {
        if (finished_ || !noMoreInput_) {
            return nullptr;
        }

        RowVectorPtr output = sortBuffer_->getOutput(maxOutputRows_);
        finished_ = (output == nullptr);
        return output;
    }

    void OrderBy::close() {
        Operator::close();
        sortBuffer_.reset();
    }
} // namespace kumo::pollux::exec
